from training.dl_frame.month_rolling.train_process_rs import Train_process

class Dl_pre_train():
    def __init__(self, args):
        self.args = args
    
    def pre_train_process(self, X_train, X_eval, Y_train,  Y_eval, times_train, times_eval,  root):
        Tr = Train_process(self.args)
        best_model = Tr.pre_train(X_train, X_eval, Y_train,  Y_eval, times_train, times_eval, root)
        train_loader,eval_loader = Tr._get_pretrain_data(X_train, X_eval, Y_train,  Y_eval, times_train, times_eval)
        train_pred, train_true = Tr.predict_future(best_model, train_loader, root, steps='pretrain_train')
        eval_pred, eval_true = Tr.predict_future(best_model, eval_loader, root, steps='pretrain_eval')
        return  train_pred, train_true, eval_pred, eval_true

    def finetune_process(self, X_ft_train, X_ft_eval, X_out, Y_ft_train,  Y_ft_eval, Y_out, times_train, times_eval, times_out, root, pre_root):
        Tr = Train_process(self.args)
        train_loader,eval_loader,predict_loader = Tr._get_finetune_data(X_ft_train, X_ft_eval, X_out, Y_ft_train,  Y_ft_eval, Y_out, times_train, times_eval, times_out)
        best_model = Tr.fintune_train(X_ft_train, X_ft_eval, X_out, Y_ft_train,  Y_ft_eval, Y_out, times_train, times_eval, times_out, root, pre_root)
        train_pred, train_true = Tr.predict_future(best_model, train_loader, root, steps='finetune_train')
        eval_pred, eval_true = Tr.predict_future(best_model, eval_loader, root, steps='finetune_eval')
        pred, true = Tr.predict_future(best_model, predict_loader, root, steps='finetune_out')
        return train_pred, train_true, eval_pred, eval_true, pred, true